import torch
import torch.nn as nn
from omod import ChinesePosParser
import pickle as pkl

with open('data/generate_pkl/depvocab.pkl', 'rb') as f:
    word_to_idx, pos_to_idx, deprel_to_idx = pkl.load(f)
    print(word_to_idx)

vocab_size = len(word_to_idx) + 1  # +1 用于未知词
pos_size = len(pos_to_idx)+1  # +1 用于未知标签
model = ChinesePosParser(vocab_size, pos_size)
model.load_state_dict(torch.load('model/pos_parser.pkl', weights_only=False))
model.eval()

w = "我 想 吃 鱼 了"
words = [[word_to_idx.get(idx, 'UNK') for idx in w.split()]]

pos_scores = model(torch.tensor(words, dtype=torch.long))
print(pos_scores)

# 获取预测结果
predicted_pos = torch.argmax(pos_scores, dim=1).tolist()

print("预测的词性:", predicted_pos)

# 假设 deprel_to_idx 是训练时使用的依存关系标签表
# deprel_to_idx = {'nsubj': 1, 'root': 2, 'dobj': 3, 'compound': 4}
idx_to_pos = {v: k for k, v in pos_to_idx.items()}

# 将索引转换为依存关系标签
predicted_pos = [idx_to_pos.get(idx, '?') for idx in predicted_pos]

# 打印结果
# for i, (word, head, deprel) in enumerate(zip(words, predicted_heads, predicted_deprels)):
#     print(f"词: {word}, 依存头: {head}, 依存关系: {deprel}")
print("词性", predicted_pos)